{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 272,
   "id": "protected-boundary",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 273,
   "id": "laden-funeral",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": 273,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kernel_dim_x = 5\n",
    "kernel_dim_y = 1 \n",
    "\n",
    "stride_x = 1\n",
    "stride_y = 1\n",
    "\n",
    "input_precision = 8\n",
    "ifm_channels = 2\n",
    "ofm_channels = 2\n",
    "ifm_dimension_x = 10\n",
    "ifm_dimension_y = 4\n",
    "padding = 0\n",
    "\n",
    "ofm_dimension_x = 6\n",
    "ofm_dimension_y = 4\n",
    "\n",
    "activation_precision = 16\n",
    "simd = 2\n",
    "pe = 2\n",
    "w_precision = 4\n",
    "mmv=2\n",
    "\n",
    "tile = ifm_channels *kernel_dim_x*kernel_dim_y * ofm_channels // (simd*pe)\n",
    "tile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 274,
   "id": "binding-thailand",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 2, 10, 4])"
      ]
     },
     "execution_count": 274,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "low = 0\n",
    "high = 4\n",
    "feature = torch.randint(low, high, (1, ifm_channels, ifm_dimension_x, ifm_dimension_y)) * 1.0\n",
    "feature.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 275,
   "id": "alternate-surname",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 2, 5, 1])"
      ]
     },
     "execution_count": 275,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv = nn.Conv2d(\n",
    "    in_channels=ifm_channels, \n",
    "    out_channels=ofm_channels, \n",
    "    kernel_size=(kernel_dim_x, kernel_dim_y), \n",
    "    stride=(stride_x, stride_y), \n",
    "    padding=padding\n",
    ")\n",
    "conv.weight = nn.Parameter(torch.randint_like(conv.weight, low, high) * 1.0)\n",
    "conv.bias = nn.Parameter(torch.zeros_like(conv.bias) * 1.0)\n",
    "w = conv.weight.detach().numpy()\n",
    "conv.weight.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 276,
   "id": "thermal-brunei",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2.0"
      ]
     },
     "execution_count": 276,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w[0,0,0,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 277,
   "id": "progressive-extreme",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 2, 6, 4])\n"
     ]
    }
   ],
   "source": [
    "print(conv(feature).shape)\n",
    "a = conv(feature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 278,
   "id": "prerequisite-champagne",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 2, 6, 4])\n"
     ]
    }
   ],
   "source": [
    "print(F.conv2d(feature, conv.weight).shape)\n",
    "b = F.conv2d(\n",
    "    feature, \n",
    "    conv.weight,\n",
    "    stride=(stride_x, stride_y),\n",
    "    padding=padding\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 279,
   "id": "american-teens",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0)"
      ]
     },
     "execution_count": 279,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(a != b).sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "worthy-objective",
   "metadata": {},
   "source": [
    "# Write to .h file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 280,
   "id": "julian-hollow",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0xfb'"
      ]
     },
     "execution_count": 280,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_hex(num, width=input_precision):\n",
    "    if num >= 0:\n",
    "        return hex(num)\n",
    "    else:\n",
    "        return hex((1 << width) + num)\n",
    "get_hex(-5, 8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 281,
   "id": "simplified-quarter",
   "metadata": {},
   "outputs": [],
   "source": [
    "outFileWeights = open(\"memdata_nonsquare.h\" , \"wt\")\n",
    "outFileConfig = open(\"config_nonsquare.h\" , \"wt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 282,
   "id": "perfect-calendar",
   "metadata": {},
   "outputs": [],
   "source": [
    "outFileConfig.write(\"#define KERNEL_DIM_X %d \\n\" % kernel_dim_x)\n",
    "outFileConfig.write(\"#define KERNEL_DIM_Y %d \\n\" % kernel_dim_y)\n",
    "outFileConfig.write(\"#define SIMD1 %d \\n\" % simd)\n",
    "outFileConfig.write(\"#define PE1 %d \\n\" % pe)\n",
    "outFileConfig.write(\"#define MMV1 %d \\n\" % mmv)\n",
    "outFileConfig.write(\"#define WIDTH %d \\n\" % w_precision)\n",
    "\n",
    "outFileConfig.write(\"#define IFM_Channels1 %d \\n\" % ifm_channels)\n",
    "outFileConfig.write(\"#define OFM_Channels1 %d \\n\" % ofm_channels)\n",
    "outFileConfig.write(\"#define IFMDim1_x %d \\n\" % ifm_dimension_x)\n",
    "outFileConfig.write(\"#define IFMDim1_y %d \\n\" % ifm_dimension_y)\n",
    "outFileConfig.write(\"#define OFMDim1_x %d \\n\" % ofm_dimension_x)\n",
    "outFileConfig.write(\"#define OFMDim1_y %d \\n\" % ofm_dimension_y)\n",
    "outFileConfig.write(\"#define STRIDE_x %d \\n\" % stride_x)\n",
    "outFileConfig.write(\"#define STRIDE_y %d \\n\" % stride_y)\n",
    "outFileConfig.write(\"#define INPUT_PRECISION %d \\n\" % input_precision)\n",
    "outFileConfig.write(\"#define TILE1 %d \\n\" % tile)\n",
    "\n",
    "outFileConfig.write(\"#define ACTIVATION_PRECISION %d \\n\" % activation_precision)\n",
    "\n",
    "outFileConfig.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 283,
   "id": "comparable-thumbnail",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0, 0, 0, 0]\n",
      "[0, 1, 0, 0]\n",
      "[0, 0, 1, 0]\n",
      "[0, 1, 1, 0]\n",
      "[0, 0, 2, 0]\n",
      "[0, 1, 2, 0]\n",
      "[0, 0, 3, 0]\n",
      "[0, 1, 3, 0]\n",
      "[0, 0, 4, 0]\n",
      "[0, 1, 4, 0]\n",
      "[1, 0, 0, 0]\n",
      "[1, 1, 0, 0]\n",
      "[1, 0, 1, 0]\n",
      "[1, 1, 1, 0]\n",
      "[1, 0, 2, 0]\n",
      "[1, 1, 2, 0]\n",
      "[1, 0, 3, 0]\n",
      "[1, 1, 3, 0]\n",
      "[1, 0, 4, 0]\n",
      "[1, 1, 4, 0]\n"
     ]
    }
   ],
   "source": [
    "outFileWeights.write(\"#ifndef PARAMS_HPP\\n\")\n",
    "outFileWeights.write(\"#define PARAMS_HPP\\n\")\n",
    "\n",
    "outFileWeights.write(\"namespace PARAM{ \\n\")\n",
    "if (w_precision == 1):\n",
    "        outFileWeights.write(\"static BinaryWeights<%d,%d,%d> weights= {\\n{\\n\" %(simd,pe,tile))\n",
    "else:\n",
    "        outFileWeights.write(\"static FixedPointWeights<%d,ap_int<%d>,%d,%d> weights= {\\n{\\n\" %(simd,w_precision,pe,tile))\n",
    "\n",
    "\n",
    "assert(ofm_channels % pe == 0)\n",
    "assert(ifm_channels % simd == 0)\n",
    "n_z = ifm_channels / simd\n",
    "n_tile = pe \n",
    "cnt = 0\n",
    "for p in range(pe):\n",
    "        outFileWeights.write(\"{ \\n\")\n",
    "        for t in range(tile):\n",
    "            val = []\n",
    "            for i in range(simd):\n",
    "                print([\n",
    "                    cnt // (kernel_dim_x*kernel_dim_y*ifm_channels), \n",
    "                    cnt % ifm_channels, \n",
    "                    (cnt // ifm_channels)%kernel_dim_x, \n",
    "                    ((cnt // ifm_channels)//kernel_dim_x)%kernel_dim_y])\n",
    "                val.append(w[\n",
    "                    cnt // (kernel_dim_x*kernel_dim_y*ifm_channels), \n",
    "                    cnt % ifm_channels, \n",
    "                    (cnt // ifm_channels)%kernel_dim_x, \n",
    "                    ((cnt // ifm_channels)//kernel_dim_x)%kernel_dim_y])\n",
    "                cnt += 1\n",
    "                pass\n",
    "            _sum = 0\n",
    "            for i, v in enumerate(val):\n",
    "                _sum += int(v) << (i * w_precision)\n",
    "            outFileWeights.write(\"%s\" % hex(_sum))\n",
    "            if t!=tile -1:\n",
    "                    outFileWeights.write(\",\\n\")\n",
    "        outFileWeights.write(\"} \\n\")\n",
    "        if p!=pe-1:\n",
    "                outFileWeights.write(\",\")\n",
    "\n",
    "\n",
    "outFileWeights.write(\"}\\n};\\n } \\n\")\n",
    "outFileWeights.write(\"#endif \\n\")\n",
    "outFileWeights.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "geographic-dimension",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
